from llama_index.llms.dashscope import DashScope, DashScopeGenerationModels
from llama_index.core.base.llms.types import MessageRole, ChatMessage
import os


USER_PROMPT = """
当前故事进展是{description}
当前故事有如下选择：{choices}
玩家选择了{user_choice}
请沿着这一选择给出进一步发展，以及该选择的三个进一步选择，你的选择务必使故事更加有趣
[返回格式]
{{
    "description": "故事后续发展。注意下面的内容不是总要说明的，偶尔你也可以在这里描述没有选择其他选项时错过的风险和收益，但注意，必须是通顺的话，因为你在写故事",
    "choices": ["三个后续选择"]
}}
你的故事需要有结局，当你认为故事可以结束时，按以下格式返回：
{{
    "description": "[END]",
    "choices": ["","",""]
}}
如果出现对话一定要转义双引号，如：
{{
    "description": "他说：\\"我一定会回来的！\\" 故事后续发展。",
    "choices": ["三个后续选择"]
}}
返回时禁止使用markdown格式
"""

SYSTEM_PROMPT = """
你是一个多分枝故事的作者，你擅长设计意外的，反转的，有趣的，引人入胜的故事，并且会使用伏笔
不要编写太长的故事，中型故事即可。在玩家耐心耗尽前要给出故事的结局。 
请对对话内容中的双引号进行转义，例如： \\"你好！\\""   
"""

class LLM:
    def __init__(self, 
                 model_name: str = DashScopeGenerationModels.QWEN_MAX,
                 system_prompt: str = "You are a helpful assistant."):
        self.llm = DashScope(model=model_name,
                             api_key=os.getenv("DASHSCOPE_API_KEY"))
        self.messages = [
            ChatMessage(
                role=MessageRole.SYSTEM, content=system_prompt
            )
        ]

    def chat(self, user_input: str):

        self.messages.append(ChatMessage(role=MessageRole.USER, content=user_input))
        resp = self.llm.chat(self.messages)
        self.messages.append(
            ChatMessage(role=MessageRole.ASSISTANT, content=resp.message.content)
        )
        return resp.message.content

    def reset(self):
        """
        重置对话历史。
        """
        self.messages = [
            ChatMessage(
                role=MessageRole.SYSTEM, content="You are a helpful assistant."
            )
        ]

